/*
 * SensorMemory.cpp
 *
 *  Created on: 2013-02-06
 *      Author: mathieu
 */

#include "semol/core/SensoryMemory.h"

#include <opencv2/opencv.hpp>

#include <utilite/ULogger.h>
#include <utilite/UColorTable.h>
#include <utilite/UMath.h>
#include <utilite/UConversion.h>

namespace semol {

SensoryMemory::SensoryMemory(const ParametersMap & parameters) :
		_audiodBThreshold(Parameters::defaultSMAudioDBThreshold()),
		_audioLog(Parameters::defaultSMAudioLog()),
		_audiodBIndexing(Parameters::defaultSMAudioDBIndexing()),
		_audioIgnoredDCComponents(Parameters::defaultSMAudioIgnoredDCComponents()),
		_audioCrop(Parameters::defaultSMAudioCrop()),
		_audioMinMagnitude(Parameters::defaultSMAudioMinMagnitude()),
		_twistDecimalResolution(Parameters::defaultSMTwistDecimalResolution()),
		_grayScale(Parameters::defaultSMImageGrayScale()),
		_lastSensorId(0)
{
	this->parseParameters(parameters);
	_sensors.resize(Sensor::kTypeEnd);
}

SensoryMemory::~SensoryMemory()
{
	clear();
}

void SensoryMemory::clear()
{
	for(unsigned int j=0; j<_sensors.size(); ++j)
	{
		if(_sensors[j].size())
		{
			UWARN("Sensory memory would be already empty here (%d sensors of type %d still in memory).", _sensors[j].size(), j);
		}
		for(std::map<int, Sensor *>::iterator i=_sensors[j].begin(); i!=_sensors[j].end(); ++i)
		{
			delete (*i).second;
		}
	}
	_sensors.clear();
	_sensorsType.clear();
	_notRefSensors.clear();
	_similaritySensorMap.clear();
	_lastSensorId = 0;

	_sensors.resize(Sensor::kTypeEnd);
}

void SensoryMemory::parseParameters(const ParametersMap & parameters)
{
	Parameters::parse(parameters, Parameters::kSMAudioDBThreshold(), _audiodBThreshold);
	Parameters::parse(parameters, Parameters::kSMAudioLog(), _audioLog);
	Parameters::parse(parameters, Parameters::kSMAudioDBIndexing(), _audiodBIndexing);
	Parameters::parse(parameters, Parameters::kSMAudioIgnoredDCComponents(), _audioIgnoredDCComponents);
	Parameters::parse(parameters, Parameters::kSMAudioCrop(), _audioCrop);
	Parameters::parse(parameters, Parameters::kSMAudioMinMagnitude(), _audioMinMagnitude);
	Parameters::parse(parameters, Parameters::kSMTwistDecimalResolution(), _twistDecimalResolution);
	Parameters::parse(parameters, Parameters::kSMImageGrayScale(), _grayScale);

}

int SensoryMemory::add(Sensor::Type type, const cv::Mat & data)
{
	UDEBUG("");
	UASSERT(type < Sensor::kTypeEnd);

	bool keepRawData = true;
	std::vector<float> descriptor;
	switch(type)
	{
		case Sensor::kTypeImage:
			descriptor = computeImageDescriptor(data);
			break;
		case Sensor::kTypeAudioFreqSqrdMagn:
			descriptor = computeAudioDescriptor(data);
			break;
		case Sensor::kTypeTwist:
			descriptor = computeTwistDescriptor(data);
			break;
		default:
			UFATAL("Not handled sensor type %d", (int)type);
			return 0;
			break;
	}

	if(descriptor.size() == 0)
	{
		UDEBUG("Bad sensor signature!");
		return 0;
	}

	int sensorId = ++_lastSensorId;
	this->add(new Sensor(sensorId, descriptor, type, keepRawData?data:cv::Mat()));

	return sensorId;
}

int SensoryMemory::add(const SensorRaw & sensor)
{
	return this->add(sensor.type(), sensor.data());
}

void SensoryMemory::add(Sensor * sensor)
{
	if(sensor)
	{
		UDEBUG("Sensor %d of type %d added to Sensory Memory", sensor->id(), sensor->type());
		UASSERT(sensor->type() < Sensor::kTypeEnd);
		if(!_sensors[sensor->type()].insert(std::make_pair(sensor->id(), sensor)).second)
		{
			UFATAL("Sensor %d already in SensoryMemory!", sensor->id());
		}
		_sensorsType.insert(std::make_pair(sensor->id(), sensor->type()));
		_notRefSensors.insert(sensor->id());
		if(sensor->id() > _lastSensorId)
		{
			_lastSensorId = sensor->id();
		}
	}
}

void SensoryMemory::add(const std::map<int, Sensor *> & sensors)
{
	for(std::map<int, Sensor *>::const_iterator iter=sensors.begin(); iter!=sensors.end(); ++iter)
	{
		this->add(iter->second);
	}
}

void SensoryMemory::addSensorRef(int sensorId, int nodeId)
{
	if(sensorId != 0)
	{
		Sensor::Type type = uValue(_sensorsType, sensorId, Sensor::kTypeEnd);
		if(type < Sensor::kTypeEnd)
		{
			Sensor * s = uValue(_sensors[type], sensorId, (Sensor*)0);
			if(s)
			{
				s->addRef(nodeId);
				_notRefSensors.erase(sensorId);
			}
			else
			{
				UERROR("Not found sensor %d", sensorId);
			}
		}
		else
		{
			UERROR("Not found sensor %d", sensorId);
		}
	}
}

void SensoryMemory::removeSensorRef(int sensorId, int nodeId)
{
	if(sensorId)
	{
		Sensor::Type type = uValue(_sensorsType, sensorId, Sensor::kTypeEnd);
		if(type < Sensor::kTypeEnd)
		{
			Sensor * s = uValue(_sensors[type], sensorId, (Sensor*)0);
			if(s)
			{
				s->removeAllRef(nodeId);
				if(s->totalReferences() == 0)
				{
					UDEBUG("Sensor %d with no more ref...", sensorId);
					_notRefSensors.insert(sensorId);
				}
			}
			else
			{
				UERROR("Not found sensor %d", sensorId);
			}
		}
		else
		{
			UERROR("Not found sensor %d", sensorId);
		}
	}
}

Sensor * SensoryMemory::remove(int id)
{
	Sensor * s=0;
	Sensor::Type type = uValue(_sensorsType, id, Sensor::kTypeEnd);
	if(type < Sensor::kTypeEnd)
	{
		s = uValue(_sensors[type], id, (Sensor*)0);
		if(s)
		{
			_sensors[type].erase(id);
		}
		_notRefSensors.erase(id);
		_sensorsType.erase(id);
	}
	return s;
}

int SensoryMemory::size() const
{
	int sum = 0;
	for(unsigned int i=0; i<_sensors.size(); ++i)
	{
		sum += _sensors[i].size();
	}
	return sum;
}

const Sensor* SensoryMemory::getSensor(int id) const
{
	Sensor::Type type = uValue(_sensorsType, id, Sensor::kTypeEnd);
	if(type < Sensor::kTypeEnd)
	{
		return uValue(_sensors[type], id, (Sensor*)0);
	}
	return 0;
}

bool SensoryMemory::contains(int id) const
{
	Sensor::Type type = uValue(_sensorsType, id, Sensor::kTypeEnd);
	if(type < Sensor::kTypeEnd)
	{
		return uContains(_sensors[type], id);
	}
	return false;
}

float SensoryMemory::compare(int sensorA, int sensorB, bool * bufferedLookup)
{
	//UDEBUG("Compare sensor %d with sensor %d", sensorA, sensorB);
	float sim = 0.0f;
	if(sensorA == 0 || sensorB == 0)
	{
		if(bufferedLookup)
		{
			*bufferedLookup = true;
		}
		if(sensorA == 0 && sensorB == 0)
		{
			//No sound should be 100% equal, because when comparing zero sound with sound, it should be 0%
			sim = 1.0f;
		}
		return sim;
	}

	if(sensorA > sensorB)
	{
		//swap to have sensorA the smallest id
		int tmp = sensorA;
		sensorA = sensorB;
		sensorB = tmp;
	}

	bool added = false;
	// Search in _similaritySensorMap <small id, big id>
	std::map<int, std::map<int, float> >::iterator iterA = _similaritySensorMap.find(sensorA);
	if(iterA != _similaritySensorMap.end())
	{
		std::map<int, float>::iterator iterAB = iterA->second.find(sensorB);
		if(iterAB != iterA->second.end())
		{
			sim = iterAB->second;
			added = true;
		}
	}
	if(bufferedLookup)
	{
		*bufferedLookup = added;
	}
	if(!added)
	{
		UASSERT_MSG(contains(sensorA), uFormat("sensor = %d", sensorA).c_str());
		UASSERT_MSG(contains(sensorB), uFormat("sensor = %d", sensorB).c_str());

		Sensor::Type typeA = uValue(_sensorsType, sensorA, Sensor::kTypeEnd);
		Sensor::Type typeB = uValue(_sensorsType, sensorB, Sensor::kTypeEnd);
		UASSERT_MSG(typeA < Sensor::kTypeEnd && typeB < Sensor::kTypeEnd && typeA == typeB, uFormat("sensorA=%d sensorB=%d typeA=%d typeB=%d", sensorA, sensorB, typeA, typeB).c_str());

		Sensor * a = uValue(_sensors[typeA], sensorA, (Sensor*)0);
		Sensor * b = uValue(_sensors[typeB], sensorB, (Sensor*)0);
		UASSERT(a && b);

		sim = this->compare(a->descriptor(), b->descriptor(), typeA);

		// Add to _similaritySensorMap
		// A to B
		if(iterA != _similaritySensorMap.end())
		{
			iterA->second.insert(std::make_pair(b->id(), sim));
		}
		else
		{
			std::map<int, float> m;
			m.insert(std::make_pair(b->id(), sim));
			_similaritySensorMap.insert(std::make_pair(a->id(), m));
		}
	}
	return sim;
}

float SensoryMemory::compare(const std::vector<float> & descriptorA, const std::vector<float> & descriptorB, Sensor::Type type) const
{
	float sim = 0.0f;
	if(descriptorA.size() == 0 || descriptorB.size() == 0)
	{
		//bad signatures...
		UWARN("bad sign");
		sim = -0.0001f; // to ignore...
	}
	else if(descriptorA.size() == descriptorB.size())
	{
		//UDEBUG("sA=%d, sB=%d, size=%d", sA->id(), sB->id(), sA->postData().size());
		float sumDiff = 0.0f;
		float notNull = 0.0f;
		for(unsigned int i=0; i<descriptorA.size(); ++i)
		{
			if(descriptorA.at(i) && descriptorB.at(i))
			{
				float d = fabs(descriptorA.at(i) - descriptorB.at(i));
				UASSERT_MSG(d >= 0.0f && d <= 1.0f, uFormat("d=%f", d).c_str());
				if(type == Sensor::kTypeImage)
				{
					if(d < 0.1)
					{
						d = 0.0f;
					}
					else
					{
						d = 1.0f;
					}
				}
				else if(type == Sensor::kTypeAudioFreqSqrdMagn)
				{
					d = 0.0f; // invariable to power
				}
				sumDiff += d;
				notNull += 1.0f;
			}
			else if(descriptorA.at(i) || descriptorB.at(i))
			{
				sumDiff += 1.0f;
				notNull += 1.0f;
			}
		}
		if(notNull > 0.0f)
		{
			sim = 1.0f - sumDiff/notNull;
		}
		else
		{
			UDEBUG("not compatible sounds, no similarity");
			sim = 1.0f;
		}
	}
	else
	{
		UERROR("Data are not the same size (%d vs %d)", (int)descriptorA.size(), (int)descriptorB.size());
		sim = -0.0001f;
	}
	return sim;
}

std::vector<float> SensoryMemory::computeImageDescriptor(const cv::Mat & imageIn)
{
	UDEBUG("");
	UASSERT(imageIn.type() == CV_8UC3 || imageIn.type() == CV_8UC1);

	//Convert to grayscale
	cv::Mat image;
	if(_grayScale)
	{
		if(imageIn.type() == CV_8UC3)
		{
			cv::cvtColor(imageIn, image, cv::COLOR_BGR2GRAY);
		}
		else
		{
			image =  imageIn;
		}
	}
	else
	{
		image = imageIn;
	}

	UDEBUG("type=%d, width=%d, height=%d, nChannels=%d, imageSize=%d,", image.type(), image.cols, image.rows, image.channels(), image.total());

	// IND transform
	int k=0;
	std::vector<float> descriptor = std::vector<float>(image.cols*image.rows*image.channels());
	for(int i=0; i<image.rows; ++i)
	{
		const cv::Mat row = image.row(i); // DON'T modify row! (it refers to const data)
		for(int j=0; j<row.cols; ++j)
		{
			if(image.channels() == 1)
			{
				descriptor[k++] = float(row.at<unsigned char>(j))/255.0f;
			}
			else
			{
				descriptor[k++] = float(row.at<unsigned char>(j*3+0))/255.0f;
				descriptor[k++] = float(row.at<unsigned char>(j*3+1))/255.0f;
				descriptor[k++] = float(row.at<unsigned char>(j*3+2))/255.0f;
			}
		}
	}
	return descriptor;
}

std::vector<float> SensoryMemory::computeAudioDescriptor(const cv::Mat & audioIn)
{
	UDEBUG("");
	UASSERT_MSG(audioIn.type() == CV_32FC1, uFormat("type=%d (CV_32FC1=%d)", audioIn.type(), CV_32FC1).c_str());

	cv::Mat audio;
	if(_audioLog)
	{
		audio = audioToLog(audioIn);
	}
	else
	{
		audio = audioIn;
	}

	int frameLength = audio.cols;
	if(_audioCrop >=0 && _audioCrop < 1.0f)
	{
		frameLength = _audioCrop * float(audio.cols);
	}
	else if(_audioCrop != 1.0f)
	{
		UWARN("_audioCrop must be between 0 and 1 (value=%f)", _audioCrop);
	}

	std::vector<float> descriptor(frameLength, 0);
	unsigned int index;
	int ignoredDC = _audioLog?0:_audioIgnoredDCComponents;
	float max;
	if(ignoredDC && audio.cols > ignoredDC)
	{
		max = uMax(((float*)audio.data) + ignoredDC, frameLength-ignoredDC, index);
		index+=ignoredDC;
	}
	else
	{
		if(ignoredDC)
		{
			UWARN("ignoredDC (%d) > frameLength (%d), may gives wrong results...", ignoredDC, frameLength);
		}
		max = uMax((float*)audio.data, frameLength, index);
	}


	UDEBUG("data.rows=%d, data.cols=%d, data.type=%d, frameCropped=%d, max=%f at %d", audio.rows, audio.cols, audio.type(), frameLength, max, index);

	bool badSignature = true;
	if(max >= _audioMinMagnitude)
	{
		///float gain = std::pow(10.0f, _audiodBGain / 10.0f);
		//max *= gain;
		// _dBThreshold < 0
		// db = 10logP, db = 20logA where P=Power and A=Amplitude
		float minLimitMagnitude = max / std::pow(10.0f, -_audiodBThreshold/10.0f);// Power ratio
		UDEBUG("max=%f minLimitMagnitude=%f _audiodBThreshold=%f", max, minLimitMagnitude, _audiodBThreshold);
		//float maxLimit = max / std::pow(10.0f, -_audiodBThreshold/20.0f);// Amplitude ratio

		for(int i=0; i<frameLength; ++i)
		{
			float val = 0.0f;
			if(i >= ignoredDC)
			{
				val = audio.at<float>(0, i);// * gain;

				if(val <= minLimitMagnitude)
				{
					val = 0.0f;
				}
				else
				{
					if(!_audiodBIndexing)
					{
						val = val / max; // set between 0 and 1
					}
					else
					{
						val = 10*std::log10(val/max);// transform to dB
						val = (val-_audiodBThreshold) / (-_audiodBThreshold); // set between 0 and 1
					}
				}
			}
			descriptor[i] = val;
			if(badSignature && descriptor[i]!=0.0f)
			{
				badSignature = false;
			}
		}
	}

	if(badSignature)
	{
		return std::vector<float>();
	}
	else
	{
		return descriptor;
	}
}

std::vector<float> SensoryMemory::computeTwistDescriptor(const cv::Mat & twist)
{
	UDEBUG("");
	UASSERT(twist.type() == CV_32FC1);
	UASSERT(twist.cols == 6);

	//don't put zero if null (not moving is a command)
	//sensors[0] = data.at<float>(0, 0)>0.01f?(int)(data.at<float>(0, 0)*100.0f):1; //x
	//sensors[1] = data.at<float>(0, 5)>0.01f?(int)(data.at<float>(0, 5)*100.0f):1; //yaw

	//different resolutions
	std::vector<float> descriptor(2);

	UASSERT(twist.at<float>(0, 0) <= 1.0f && twist.at<float>(0, 0) >= -1.0f);
	UASSERT(twist.at<float>(0, 5) <= 1.0f && twist.at<float>(0, 5) >= -1.0f);

	// set between -0.5 and 0.5
	descriptor[0] = twist.at<float>(0, 0)/2.0f; //x [-1 1]
	descriptor[1] = twist.at<float>(0, 5)/2.0f; //yaw [-1 1]

	UDEBUG("twist x=%f yaw=%f, descriptor=[%f, %f]", twist.at<float>(0, 0), twist.at<float>(0, 5), descriptor[0], descriptor[1]);

	return descriptor;
}

cv::Mat SensoryMemory::audioToLog(const cv::Mat & audio)
{
	UASSERT_MSG(audio.type() == CV_32FC1 && audio.cols > 0, uFormat("type=%d (CV_32FC1=%d)", audio.type(), CV_32FC1).c_str());

	// Log stuff
	int nbSubOctave = 16;
	cv::Mat dataLog(1, log2(audio.cols) * nbSubOctave + 1, CV_32FC1);

	int octaveStart = 1; // DC ignored
	float avgMag = 0;
	int count = 0;
	int j = 0; // dataLog index
	int k = 0; //
	float subOctaveWidth = float(octaveStart)/float(nbSubOctave);
	int ignoredDC = _audioIgnoredDCComponents;
	for(int i=0; i<audio.cols; ++i)
	{
		if(i>=octaveStart)
		{
			// log indexes = 1 2 4 8 16 32 64 128 256 512
			// sub2-log indexes 1 1.5 2 3 4 6 8 12 16 24 32 48 64 96 128 192 256 384 512
			// sub3-log indexes 1 1.3 1.6 2...
			avgMag += i>=ignoredDC?audio.at<float>(0,i):0;
			++count;
			if(i >= octaveStart + int(float(k+1)*subOctaveWidth) - 1)
			{
				avgMag/=float(count);
				//UDEBUG("i=%d, j=%d, k=%d octaveStart=%d (suboct=%d) -> avg=%f, count=%d, octWidth=%f", i, j, k, octaveStart, octaveStart + int(float(k)*subOctaveWidth), avgMag, count, subOctaveWidth);
				dataLog.at<float>(0,j++) = avgMag;
				while(i > octaveStart + int(float(k+1)*subOctaveWidth) - 1)
				{
					++k;
					//UDEBUG("i=%d, j=%d, k=%d octaveStart=%d (suboct=%d) -> avg=%f, count=%d, octWidth=%f", i, j, k, octaveStart, octaveStart + int(float(k)*subOctaveWidth), avgMag, count, subOctaveWidth);
					dataLog.at<float>(0,j++) = avgMag;
				}
				avgMag = 0;
				count = 0;
				if(++k >= nbSubOctave)
				{
					k = 0;
					octaveStart *= 2;
					subOctaveWidth = float(octaveStart)/float(nbSubOctave);
				}
			}
		}
	}
	if(dataLog.cols > j)
	{
		UDEBUG("dataLog = %d j=%d", dataLog.cols, j);
		dataLog = cv::Mat(dataLog, cv::Range::all(), cv::Range(0,j));
	}

	return dataLog;
}

} /* namespace semol */
